from ..dlgo.board import Move, GameState
from ..dlgo.utils import Point, Player
import numpy as np 

class ZeroEncoder:
    def __init__(self, board_size=19) -> None:
        self.board_size = board_size

        # 0 - 3. our stones with 1, 2, 3, 4+ liberties
        # 4 - 7. opponent stones with 1, 2, 3, 4+ liberties
        # 8. 1 if we get komi
        # 9. 1 if opponent gets komi
        # 10. move would be illegal due to ko
        self.num_planes = 11

    def Encode(self, game_state:GameState):
        board_tensor = np.zeros(self.Shape())
        next_player = game_state.next_player
        if game_state.next_player == Player.white:
            board_tensor[8] = 1
        else:
            board_tensor[9] = 1
        for r in range(self.board_size):
            for c in range(self.board_size):
                p = Point(r+1, c+1)
                gostrng = game_state.board.GetGoString(p)
                if gostrng is None:
                    if game_state.DoesMoveViolateKo(next_player, Move.Play(p)):
                        board_tensor[10][r][c] = 1
                else:
                    liberty_plane = min(4, gostrng.NumberLiberties) - 1
                    if gostrng.next_player != next_player:
                        liberty_plane += 4
                    board_tensor[liberty_plane][r][c] = 1
        return board_tensor

    def EncodeMove(self, move:Move):
        if move.is_play:
            return (self.board_size * (move.point.row - 1) + (move.point.col - 1))
        elif move.is_pass:
            return self.board_size * self.board_size
        raise ValueError('Error In ZeroEncoder.EncodeMove() 36 line: Cannot encode resigne move')

    def DecodeMove(self, index:int):
        if index == self.board_size * self.board_size:
            return Move.Pass()
        r = index // self.board_size
        c = index % self.board_size
        return Move.Play(Point(r+1, c+1))

    def NumberMoves(self):
        return self.board_size * self.board_size + 1

    def Shape(self):
        return self.num_planes, self.board_size, self.board_size
